{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "057f6e78",
   "metadata": {},
   "source": [
    "# 数据下载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "18511b81",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['trans.ipynb',\n",
       " 'README.md',\n",
       " 'MNIST_data',\n",
       " '.ipynb_checkpoints',\n",
       " 'fashion.ipynb',\n",
       " 'FashionMNIST',\n",
       " 't.sh',\n",
       " 'transformer_learn.py']"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "os.listdir()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b5e3ccb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "from torchvision import datasets, transforms\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "370634b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 对图片数据进行转换，因为Fashion minist是灰度图，所以使用下面的normalization方式\n",
    "transform = transforms.Compose([\n",
    "transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2427b554",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 数据下载与封装\n",
    "train_set = datasets.FashionMNIST('./', download = True, train = True, transform = transform)\n",
    "test_set = datasets.FashionMNIST('./', download = True, train = False, transform = transform)\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size = 64, shuffle = True)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size = 64, shuffle = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e629e9c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b7e549c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.Tensor'>\n",
      "torch.Size([64, 1, 28, 28])\n",
      "torch.Size([64])\n"
     ]
    }
   ],
   "source": [
    "# 检查数据\n",
    "dataiter = iter(train_loader)\n",
    "images, labels = dataiter.next()\n",
    "print(type(images))\n",
    "print(images.shape)\n",
    "print(labels.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cb66c232",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7fcb6abe3700>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAD4CAYAAAAq5pAIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4xLjAsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+17YcXAAAR60lEQVR4nO3db4yV5ZkG8OsSBkEGx2FAFhiUlmBEROhK0MgGWP+gYCLWpJvyoXETstMPbdKaflijJvVLE7OxrU3UJtPVlK6utUlLwMTEEmxCmhh0RMrgojiLbh0ZGP7qAMrfez/MazPVee/7eN5zznvguX4JmZlzzzvnmQMX7znnfp/noZlBRC5+l5Q9ABFpDIVdJBEKu0giFHaRRCjsIokY28g7I6m3/uugvb09t9ba2uoeO2HCBLcedWvOnj3r1k+ePJlbO3jwoHvs+fPn3bqMzsw42u2Fwk7yLgC/ADAGwH+a2WNFfp5U5/bbb8+tLV++3D123rx5bv3MmTNu/ejRo279rbfeyq099dRT7rEnTpxw6/LVVP00nuQYAE8BWAXgOgBrSV5Xq4GJSG0Vec2+BECfme01s9MAfgtgTW2GJSK1ViTsMwF8OOLr/uy2v0Oyi2QPyZ4C9yUiBRV5zT7amwBfejfHzLoBdAN6g06kTEXO7P0AZo34uhPAvmLDEZF6KRL2NwDMJfk1kuMAfBvAptoMS0RqjUVmvZFcDeAJDLfenjWznwTff1E+jSdHbWv+TfQY33rrrW796aefduvjx4/PrbW0tLjHRq21yLlz59z62LH5rxSjxyWq33jjjW79yJEjubUxY8a4x0a/VzOrS5/dzF4G8HKRnyEijaHLZUUSobCLJEJhF0mEwi6SCIVdJBEKu0giGjqf/WLl9ZKBuJe9ZMkSt37FFVe4dW8qaNQvjvrNkegaA68ePS4TJ05066tWrXLrzz//vFtPjc7sIolQ2EUSobCLJEJhF0mEwi6SCIVdJBFqvdVA0WmifX19bj1aUtmbxrpx40b32KVLl7r1aKnpaKnqV155Jbd20003ucdeeumlbj163DwX8hTWaunMLpIIhV0kEQq7SCIUdpFEKOwiiVDYRRKhsIskQn32BoimYj7wwANuPZpG6vWMly1b5h770ksvufWuri63PjQ05Na95aDb2trcY0+fPu3W161b59Z7e3tza95W0hcrndlFEqGwiyRCYRdJhMIukgiFXSQRCrtIIhR2kUSoz14Dr7/+uluPlkSOetWHDx9265MmTcqtzZkzxz32mmuucev9/f1u/ezZs2591qxZubVoGetDhw659VtuucWt79u3L7f2yCOPuMc++eSTbv1CVCjsJD8AMATgHICzZra4FoMSkdqrxZn9n83M/y9YREqn1+wiiSgadgPwR5Jvkhz1ImqSXSR7SPYUvC8RKaDo0/ilZraP5JUANpN8x8y2jvwGM+sG0A0AJPNnRYhIXRU6s5vZvuzjIIANAPwdCkWkNFWHneREkpM+/xzASgC7ajUwEaktevON3QPJr2P4bA4Mvxz4bzP7SXDMBfs0/p577smtPf744+6xg4ODbj3qw48fP96te+vKjxs3zj02+vt/8cUX3fptt93m1r1rBObPn+8eG805P378uFv31gGIfu+bb77ZrTczMxv1F6/6NbuZ7QWwsOoRiUhDqfUmkgiFXSQRCrtIIhR2kUQo7CKJ0BTXCt133325tUsu8f/P9KagAvHWxBFvmmm0HHO0nPPKlSvd+v79+92611astu37uWg76c8++yy3Fv3e0RLcW7dudevNSGd2kUQo7CKJUNhFEqGwiyRCYRdJhMIukgiFXSQR6rNX6Prrr8+tef1cIO4nR8dHyzV7U2CjKa4ff/yxW+/s7CxUP3r0aG4t+r3PnDnj1qPjvd89WsZ69erVbl19dhFpWgq7SCIUdpFEKOwiiVDYRRKhsIskQmEXSYT67BXq6OjIrUVbC1922WVu3VsKGojnu3tz1r3llIF4GesPP/zQrU+ePNmtnzhxIrcW9eg/+eQTtx5df+CtM3Dq1Cn32IULL76Fk3VmF0mEwi6SCIVdJBEKu0giFHaRRCjsIolQ2EUSoT57prW11a1HvXJPtK58NOc8Wvvd62VHveixY/1/AtG87+gaA29t9wMHDrjHHjt2zK1PmzbNrXvz/D/99FP32KlTp7r1C1F4Zif5LMlBkrtG3DaZ5GaS72Uf2+s7TBEpqpKn8b8GcNcXbnsQwBYzmwtgS/a1iDSxMOxmthXAkS/cvAbA+uzz9QDurfG4RKTGqn3NPs3MBgDAzAZIXpn3jSS7AHRVeT8iUiN1f4POzLoBdAMAyWI7+YlI1aptvR0gOR0Aso+DtRuSiNRDtWHfBOD+7PP7AWyszXBEpF7Cp/EkXwCwAsAUkv0AfgzgMQC/I7kOwF8BfKueg2yEa6+91q17a7+fO3fOPTbaRzyaU97f3+/WvbFFa9ZHY49EfXpv7feohx/Vo2sfvOsXjh8/7h4bzdO/EIVhN7O1OaXbajwWEakjXS4rkgiFXSQRCrtIIhR2kUQo7CKJ0BTXzLx589y6tz1wNEU1WrY4mgIbtYG89lnUGotab9HYoqWqvem50TTTyy+/3K2/++67bn3mzJm5tej3amlpcetRW7BoS7MedGYXSYTCLpIIhV0kEQq7SCIUdpFEKOwiiVDYRRKhPntmxowZbt3rm0ZTWDdt2uTWly9f7ta9fjHgLxcdbQcdTYGNlrGeNGmSW/e2mz58+LB7bHR9wYYNG9z6smXLcmvRYxotkT1nzhy3vmfPHrdeBp3ZRRKhsIskQmEXSYTCLpIIhV0kEQq7SCIUdpFEqM+eaWtrc+vefPYpU6a4x+7cudOtL1q0yK1fddVVbt3rs3vbFgPxXPxoyWVvu2jA7+NHc8qj6xe2b9/u1qdPn55b6+zsdI+N5qsvWLDAravPLiKlUdhFEqGwiyRCYRdJhMIukgiFXSQRCrtIItRnz0RrlHt99mht9m3btrn1G264wa0vWbLErXtbF0frukd9+KiPHh3vzYeP1laPfvbBgwervu/oZ3t/35Uc34zCMzvJZ0kOktw14rZHSX5Eckf2Z3V9hykiRVXyNP7XAO4a5fafm9mi7M/LtR2WiNRaGHYz2wrgSAPGIiJ1VOQNuu+T3Jk9zW/P+yaSXSR7SPYUuC8RKajasP8SwBwAiwAMAPhp3jeaWbeZLTazxVXel4jUQFVhN7MDZnbOzM4D+BUA/+1iESldVWEnOXLu4DcB7Mr7XhFpDmGfneQLAFYAmEKyH8CPAawguQiAAfgAwHfrOMaGiNYR9+ZlR73saG7z4OCgW4/mVnvzws+cOeMeG63dHvWTo73n29tz384J58pHffjocfPq0Xr50d9ptH97MwrDbmZrR7n5mTqMRUTqSJfLiiRCYRdJhMIukgiFXSQRCrtIIjTFtUJei8nblrgS3hRVwF8qGgBOnjyZW4taZ1GLKWqtRS0sb4ps9HtF9ehxP3Ikf0pH9LOjra6jJbibkc7sIolQ2EUSobCLJEJhF0mEwi6SCIVdJBEKu0gi1GfPREtJF5lGGomm10a98Gjr4yKi6bXRMtreY1O0hz9//ny3PjQ0lFvr6Ohwj42m30bbSTcjndlFEqGwiyRCYRdJhMIukgiFXSQRCrtIIhR2kUSoz54pMu872t43Em3ZHG2b7PW6ox59JLqGIOrDe3P1vXn4ldQXLlzo1nfv3p1bi3r8kba2tkLHl0FndpFEKOwiiVDYRRKhsIskQmEXSYTCLpIIhV0kEeqzZ6J1xL1+cdRnj3rR0Xz2o0ePunWvFx7ddzQvO9o2+fTp027dG1t031GPf+7cuW79tddey61F11VELsp140nOIvknkrtJvk3yB9ntk0luJvle9jF/I24RKV0lT+PPAviRmc0DcDOA75G8DsCDALaY2VwAW7KvRaRJhWE3swEz2559PgRgN4CZANYAWJ9923oA99ZrkCJS3Fd6zU5yNoBvANgGYJqZDQDD/yGQvDLnmC4AXcWGKSJFVRx2kq0Afg/gh2b2SaUTLMysG0B39jP8FQRFpG4qar2RbMFw0J83sz9kNx8gOT2rTwcwWJ8hikgthGd2Dp/CnwGw28x+NqK0CcD9AB7LPm6sywgbJFq22JtGun//fvfYGTNmuPVoGeuo9dba2ppbi9pX0XbR0fHRMzyvNRe1O6Mlsjs7O916X19fbq3oNtvR2JtRJU/jlwL4DoBekjuy2x7CcMh/R3IdgL8C+FZ9higitRCG3cz+DCDvv+/bajscEakXXS4rkgiFXSQRCrtIIhR2kUQo7CKJ0BTXTNQv9qZEfvTRR+6xd999t1uPerZRr9vrs0fTb48dO+bWo8fl/Pnzbt0bWzR9tsj1BYDfS4+2ZC66VXUz0pldJBEKu0giFHaRRCjsIolQ2EUSobCLJEJhF0nEhdcsrJNo7nQ0391zxx13uPWoFx712YeGhnJr0ZLHUa87EvWbvW2Xo8c8Glu0DsDVV1+dW4se8+hxi/rwzUhndpFEKOwiiVDYRRKhsIskQmEXSYTCLpIIhV0kEeqzZw4dOuTWFyxYkFt77rnn3GPb2/0NblesWOHWozXOvZ5wNB89+tlFri+Iji+yVj8AtLW1ufX3338/t7Z371732NmzZ7v1PXv2uPVmpDO7SCIUdpFEKOwiiVDYRRKhsIskQmEXSYTCLpKISvZnnwXgNwD+AcB5AN1m9guSjwL4NwAHs299yMxertdA6623t9et33nnnbm1aH/2mTNnuvVobnW0NvupU6dya1EvO5pTHonG5v38aGzRevrRfbe0tOTWvDUAKvnZF2KfvZKLas4C+JGZbSc5CcCbJDdntZ+b2eP1G56I1Eol+7MPABjIPh8iuRuAf6oSkabzlZ7DkZwN4BsAtmU3fZ/kTpLPkhz1mlCSXSR7SPYUGqmIFFJx2Em2Avg9gB+a2ScAfglgDoBFGD7z/3S048ys28wWm9niGoxXRKpUUdhJtmA46M+b2R8AwMwOmNk5MzsP4FcAltRvmCJSVBh2Dk+begbAbjP72Yjbp4/4tm8C2FX74YlIrVTybvxSAN8B0EtyR3bbQwDWklwEwAB8AOC7dRlhg0TTUCdMmJBbi5Yd9pY0BoCpU6e69RMnTrh1b0nlolNUi/LaZ9H026j11tHRUXU9OjaaXhttF92MKnk3/s8ARvtbuWB76iIp0hV0IolQ2EUSobCLJEJhF0mEwi6SCIVdJBFaSjrT0+Nfuv/OO+/k1gYGBtxjH374Ybf+xBNPuPUdO3a4dW9Z5Gip6Ggq5/jx4926N40U8Hvl0TUAUX3atGlu/dVXX82t9fX1FfrZ0fHNSGd2kUQo7CKJUNhFEqGwiyRCYRdJhMIukgiFXSQRbOR8Z5IHAfzfiJumAPD3Si5Ps46tWccFaGzVquXYrjazURdIaGjYv3TnZE+zrk3XrGNr1nEBGlu1GjU2PY0XSYTCLpKIssPeXfL9e5p1bM06LkBjq1ZDxlbqa3YRaZyyz+wi0iAKu0giSgk7ybtIvkuyj+SDZYwhD8kPSPaS3FH2/nTZHnqDJHeNuG0yyc0k38s++gveN3Zsj5L8KHvsdpBcXdLYZpH8E8ndJN8m+YPs9lIfO2dcDXncGv6aneQYAHsA3AGgH8AbANaa2f80dCA5SH4AYLGZlX4BBsllAI4D+I2ZXZ/d9h8AjpjZY9l/lO1m9u9NMrZHARwvexvvbLei6SO3GQdwL4B/RYmPnTOuf0EDHrcyzuxLAPSZ2V4zOw3gtwDWlDCOpmdmWwEc+cLNawCszz5fj+F/LA2XM7amYGYDZrY9+3wIwOfbjJf62Dnjaogywj4TwIcjvu5Hc+33bgD+SPJNkl1lD2YU08xsABj+xwPgypLH80XhNt6N9IVtxpvmsatm+/Oiygj7aFtJNVP/b6mZ/SOAVQC+lz1dlcpUtI13o4yyzXhTqHb786LKCHs/gFkjvu4EsK+EcYzKzPZlHwcBbEDzbUV94PMddLOPgyWP52+aaRvv0bYZRxM8dmVuf15G2N8AMJfk10iOA/BtAJtKGMeXkJyYvXECkhMBrETzbUW9CcD92ef3A9hY4lj+TrNs4523zThKfuxK3/7czBr+B8BqDL8j/78AHi5jDDnj+jqAv2R/3i57bABewPDTujMYfka0DkAHgC0A3ss+Tm6isf0XgF4AOzEcrOklje2fMPzScCeAHdmf1WU/ds64GvK46XJZkUToCjqRRCjsIolQ2EUSobCLJEJhF0mEwi6SCIVdJBH/D3WO9HRX2M6KAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(images[5].numpy().squeeze(), cmap = 'Greys_r')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2f11180",
   "metadata": {},
   "source": [
    "# DNN神经网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "56d83f5c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1/30.. Training loss: 0.597.. Test loss: 0.450.. Test Accuracy: 0.831\n",
      "Epoch: 2/30.. Training loss: 0.453.. Test loss: 0.442.. Test Accuracy: 0.839\n",
      "Epoch: 3/30.. Training loss: 0.419.. Test loss: 0.401.. Test Accuracy: 0.856\n",
      "Epoch: 4/30.. Training loss: 0.398.. Test loss: 0.374.. Test Accuracy: 0.867\n",
      "Epoch: 5/30.. Training loss: 0.389.. Test loss: 0.375.. Test Accuracy: 0.868\n",
      "Epoch: 6/30.. Training loss: 0.373.. Test loss: 0.396.. Test Accuracy: 0.855\n",
      "Epoch: 7/30.. Training loss: 0.365.. Test loss: 0.385.. Test Accuracy: 0.861\n",
      "Epoch: 8/30.. Training loss: 0.358.. Test loss: 0.370.. Test Accuracy: 0.869\n",
      "Epoch: 9/30.. Training loss: 0.349.. Test loss: 0.358.. Test Accuracy: 0.877\n",
      "Epoch: 10/30.. Training loss: 0.341.. Test loss: 0.363.. Test Accuracy: 0.876\n",
      "Epoch: 11/30.. Training loss: 0.335.. Test loss: 0.360.. Test Accuracy: 0.872\n",
      "Epoch: 16/30.. Training loss: 0.316.. Test loss: 0.356.. Test Accuracy: 0.876\n",
      "Epoch: 17/30.. Training loss: 0.314.. Test loss: 0.369.. Test Accuracy: 0.872\n",
      "Epoch: 18/30.. Training loss: 0.311.. Test loss: 0.343.. Test Accuracy: 0.884\n",
      "Epoch: 19/30.. Training loss: 0.309.. Test loss: 0.338.. Test Accuracy: 0.881\n",
      "Epoch: 20/30.. Training loss: 0.305.. Test loss: 0.338.. Test Accuracy: 0.884\n",
      "Epoch: 21/30.. Training loss: 0.304.. Test loss: 0.352.. Test Accuracy: 0.877\n",
      "Epoch: 22/30.. Training loss: 0.298.. Test loss: 0.357.. Test Accuracy: 0.881\n",
      "Epoch: 23/30.. Training loss: 0.300.. Test loss: 0.349.. Test Accuracy: 0.880\n",
      "Epoch: 24/30.. Training loss: 0.294.. Test loss: 0.361.. Test Accuracy: 0.878\n",
      "Epoch: 25/30.. Training loss: 0.296.. Test loss: 0.346.. Test Accuracy: 0.884\n",
      "Epoch: 26/30.. Training loss: 0.286.. Test loss: 0.357.. Test Accuracy: 0.882\n",
      "Epoch: 27/30.. Training loss: 0.291.. Test loss: 0.360.. Test Accuracy: 0.876\n",
      "Epoch: 28/30.. Training loss: 0.290.. Test loss: 0.343.. Test Accuracy: 0.886\n",
      "Epoch: 29/30.. Training loss: 0.289.. Test loss: 0.345.. Test Accuracy: 0.884\n",
      "Epoch: 30/30.. Training loss: 0.281.. Test loss: 0.339.. Test Accuracy: 0.887\n"
     ]
    }
   ],
   "source": [
    "# Define the network architecture\n",
    "from torch import nn, optim\n",
    "import torch.nn.functional as F\n",
    "\n",
    "model = nn.Sequential(nn.Linear(784, 256),\n",
    "                      nn.ReLU(),\n",
    "                      nn.Dropout(0.2),\n",
    "                      nn.Linear(256, 128),\n",
    "                      nn.ReLU(),\n",
    "                      nn.Dropout(0.2),\n",
    "                      nn.Linear(128, 64),\n",
    "                      nn.ReLU(),\n",
    "                      nn.Dropout(0.2),\n",
    "                      nn.Linear(64, 10),\n",
    "                      nn.LogSoftmax(dim = 1)\n",
    "                     )\n",
    "\n",
    "# Define the loss\n",
    "# NLLLoss 是似然损失函数的一种，参考：https://zhuanlan.zhihu.com/p/83283586\n",
    "# 常用在分类任务中\n",
    "criterion = nn.NLLLoss()\n",
    "\n",
    "# Define the optimizer\n",
    "optimizer = optim.Adam(model.parameters(), lr = 0.002)\n",
    "\n",
    "# Define the epochs\n",
    "epochs = 30\n",
    "\n",
    "train_losses, test_losses = [], []\n",
    "for e in range(epochs):\n",
    "    running_loss = 0\n",
    "    model.train()\n",
    "    for images, labels in train_loader:\n",
    "    # Flatten Fashion-MNIST images into a 784 long vector\n",
    "        images = images.view(images.shape[0], -1)\n",
    "        \n",
    "        # Training pass\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model.forward(images)\n",
    "        loss = criterion(output, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "    test_loss = 0\n",
    "    accuracy = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        for images, labels in test_loader:\n",
    "            images = images.view(images.shape[0], -1)\n",
    "            log_ps = model(images)\n",
    "            test_loss += criterion(log_ps, labels)\n",
    "            \n",
    "            ps = torch.exp(log_ps)\n",
    "            top_p, top_class = ps.topk(1, dim = 1)\n",
    "            equals = top_class == labels.view(*top_class.shape)\n",
    "            accuracy += torch.mean(equals.type(torch.FloatTensor))\n",
    "    \n",
    "    train_losses.append(running_loss/len(train_loader))\n",
    "    test_losses.append(test_loss/len(test_loader))\n",
    "    \n",
    "    print(\"Epoch: {}/{}..\".format(e+1, epochs),\n",
    "          \"Training loss: {:.3f}..\".format(running_loss/len(train_loader)),\n",
    "          \"Test loss: {:.3f}..\".format(test_loss/len(test_loader)),\n",
    "          \"Test Accuracy: {:.3f}\".format(accuracy/len(test_loader)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4c18b5d",
   "metadata": {},
   "source": [
    "# Resnet神经网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9db47976",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "649751cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = torchvision.models.resnet18(pretrained=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "85645a12",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO:这里的7,7尺寸是什么？\n",
    "for name, layer in model.named_modules():\n",
    "    if isinstance(layer, nn.MaxPool2d):\n",
    "        model.maxpool = nn.AdaptiveAvgPool2d((7, 7))\n",
    "# 因为数据集是10分类，所以这里设置为10\n",
    "model.fc.out_features = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f6ed20d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 因为resnet不进行修改的话，只能用于彩色图片，所以这里要对数据进行处理\n",
    "# 我们可以在前面加个卷积核，将通道映射到3，然后输入到resnet中, 也可以直接修改resnet第一层通道为1\n",
    "model.conv1 = nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "852b425b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ResNet(\n",
      "  (conv1): Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3))\n",
      "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "  (relu): ReLU(inplace=True)\n",
      "  (maxpool): AdaptiveAvgPool2d(output_size=(7, 7))\n",
      "  (layer1): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer2): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer3): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (layer4): Sequential(\n",
      "    (0): BasicBlock(\n",
      "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (downsample): Sequential(\n",
      "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
      "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      )\n",
      "    )\n",
      "    (1): BasicBlock(\n",
      "      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "      (relu): ReLU(inplace=True)\n",
      "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
      "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
      "    )\n",
      "  )\n",
      "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
      "  (fc): Linear(in_features=512, out_features=10, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# 检查模型是否可以自适应了，发现已经替换完成\n",
    "# 查看模型, 因为resnet模型对输入的图片尺寸有要求，所以这里要改成自适应\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "e1f5c844",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.81it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.93it/s]\n",
      "  0%|          | 4/938 [00:00<00:30, 31.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 0, train loss:0.3041, train acc:0.8911\n",
      "test loss:0.3176, test acc:0.8881\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.80it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.49it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 1, train loss:0.2453, train acc:0.9100\n",
      "test loss:0.2984, test acc:0.8921\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.85it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.93it/s]\n",
      "  0%|          | 3/938 [00:00<00:31, 29.96it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 2, train loss:0.2545, train acc:0.9079\n",
      "test loss:0.2772, test acc:0.9005\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.94it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.47it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.53it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 3, train loss:0.2067, train acc:0.9239\n",
      "test loss:0.2790, test acc:0.9062\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.76it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.32it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 4, train loss:0.1901, train acc:0.9300\n",
      "test loss:0.2726, test acc:0.9089\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.83it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.15it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.60it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 5, train loss:0.1629, train acc:0.9398\n",
      "test loss:0.2578, test acc:0.9101\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.81it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.14it/s]\n",
      "  0%|          | 0/938 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 6, train loss:0.1351, train acc:0.9497\n",
      "test loss:0.2674, test acc:0.9123\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.71it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.92it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.89it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 7, train loss:0.1116, train acc:0.9574\n",
      "test loss:0.2770, test acc:0.9104\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.78it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.28it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.75it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 8, train loss:0.0979, train acc:0.9629\n",
      "test loss:0.3160, test acc:0.8997\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.71it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.96it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 9, train loss:0.0835, train acc:0.9688\n",
      "test loss:0.3152, test acc:0.9152\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.78it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.31it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:10, train loss:0.0583, train acc:0.9784\n",
      "test loss:0.3251, test acc:0.9165\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.69it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.95it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:11, train loss:0.0471, train acc:0.9827\n",
      "test loss:0.3519, test acc:0.9183\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.80it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.30it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:12, train loss:0.0362, train acc:0.9870\n",
      "test loss:0.3867, test acc:0.9177\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.76it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.24it/s]\n",
      "  0%|          | 3/938 [00:00<00:31, 29.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:13, train loss:0.0315, train acc:0.9886\n",
      "test loss:0.4099, test acc:0.9190\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.77it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.95it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.17it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:14, train loss:0.0212, train acc:0.9924\n",
      "test loss:0.4525, test acc:0.9168\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.82it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.33it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:15, train loss:0.0169, train acc:0.9941\n",
      "test loss:0.4839, test acc:0.9173\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.29it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 66.36it/s]\n",
      "  0%|          | 3/938 [00:00<00:31, 29.31it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:16, train loss:0.0143, train acc:0.9948\n",
      "test loss:0.4798, test acc:0.9161\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:28<00:00, 33.17it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.55it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:17, train loss:0.0094, train acc:0.9967\n",
      "test loss:0.5085, test acc:0.9200\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.51it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.18it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:18, train loss:0.0080, train acc:0.9971\n",
      "test loss:0.5224, test acc:0.9190\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:27<00:00, 33.63it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.42it/s]\n",
      "  0%|          | 3/938 [00:00<00:35, 26.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:19, train loss:0.0066, train acc:0.9978\n",
      "test loss:0.5896, test acc:0.9194\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:30<00:00, 31.13it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 66.58it/s]\n",
      "  0%|          | 3/938 [00:00<00:33, 27.74it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:20, train loss:0.0061, train acc:0.9983\n",
      "test loss:0.6530, test acc:0.9191\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.31it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 67.32it/s]\n",
      "  0%|          | 3/938 [00:00<00:32, 28.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:21, train loss:0.0044, train acc:0.9986\n",
      "test loss:0.6422, test acc:0.9179\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.58it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 67.33it/s]\n",
      "  0%|          | 3/938 [00:00<00:32, 28.63it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:22, train loss:0.0018, train acc:0.9995\n",
      "test loss:0.6614, test acc:0.9195\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.69it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 66.69it/s]\n",
      "  0%|          | 3/938 [00:00<00:32, 28.35it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:23, train loss:0.0036, train acc:0.9990\n",
      "test loss:0.6434, test acc:0.9183\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.47it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 67.33it/s]\n",
      "  0%|          | 3/938 [00:00<00:32, 28.48it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:24, train loss:0.0025, train acc:0.9993\n",
      "test loss:0.6562, test acc:0.9206\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.49it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 67.00it/s]\n",
      "  0%|          | 3/938 [00:00<00:32, 28.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:25, train loss:0.0015, train acc:0.9996\n",
      "test loss:0.6939, test acc:0.9214\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.54it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 66.80it/s]\n",
      "  0%|          | 3/938 [00:00<00:33, 27.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:26, train loss:0.0018, train acc:0.9994\n",
      "test loss:0.6994, test acc:0.9210\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:28<00:00, 32.97it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.95it/s]\n",
      "  0%|          | 4/938 [00:00<00:29, 31.80it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:27, train loss:0.0011, train acc:0.9997\n",
      "test loss:0.7040, test acc:0.9184\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:29<00:00, 31.68it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 67.35it/s]\n",
      "  0%|          | 3/938 [00:00<00:33, 28.21it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:28, train loss:0.0009, train acc:0.9997\n",
      "test loss:0.7510, test acc:0.9215\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:28<00:00, 33.32it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.67it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:29, train loss:0.0005, train acc:0.9998\n",
      "test loss:0.7725, test acc:0.9203\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "batch_size = 128\n",
    "epochs = 30\n",
    "lr = 0.001\n",
    "gamma = 0.9\n",
    "step_size=5\n",
    "\n",
    "device_count = torch.cuda.device_count()\n",
    "USE_CUDA = torch.cuda.is_available()\n",
    "device = torch.device(\"cuda:0\" if USE_CUDA else \"cpu\")\n",
    "\n",
    "if device_count > 1:\n",
    "    model = nn.DataParallel(model,device_ids=range(device_count)) # multi-GPU\n",
    "    model.to(device)\n",
    "\n",
    "else:\n",
    "    model = model.cuda()\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma) # 学习方式\n",
    "\n",
    "# 训练模型\n",
    "train_acc, test_acc = [], []\n",
    "for epoch in range(epochs):\n",
    "    epoch_loss = 0\n",
    "    epoch_accuracy = 0\n",
    "    for data, label in tqdm(train_loader):\n",
    "        data = data.cuda()\n",
    "        label = label.cuda()\n",
    "        output = model(data)\n",
    "    \n",
    "        loss = criterion(output, label)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        acc = (output.argmax(dim=1) == label).float().mean()\n",
    "        epoch_accuracy += acc / len(train_loader)\n",
    "        epoch_loss += loss / len(train_loader)\n",
    "        \n",
    "        data.cpu()\n",
    "        label.cpu()\n",
    "        \n",
    "    with torch.no_grad():\n",
    "        epoch_test_accuracy = 0\n",
    "        epoch_test_loss = 0\n",
    "        for data, label in tqdm(test_loader):\n",
    "            data = data.cuda()\n",
    "            label = label.cuda()\n",
    "            \n",
    "            test_output = model(data)\n",
    "            test_loss = criterion(test_output, label)\n",
    "            \n",
    "            acc = (test_output.argmax(dim=1) == label).float().mean()\n",
    "            epoch_test_accuracy += acc / len(test_loader)\n",
    "            epoch_test_loss += test_loss / len(test_loader)\n",
    "            data.cpu()\n",
    "            label.cpu()\n",
    "    scheduler.step()\n",
    "    \n",
    "    print(f'EPOCH:{epoch:2}, train loss:{epoch_loss:.4f}, train acc:{epoch_accuracy:.4f}')\n",
    "    print(f'test loss:{epoch_test_loss:.4f}, test acc:{epoch_test_accuracy:.4f}')\n",
    "\n",
    "    train_acc.append(epoch_accuracy)\n",
    "    test_acc.append(epoch_test_accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92fe0df0",
   "metadata": {},
   "source": [
    "# LSTM神经网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb2b1e3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LSTM做图像分类的原理是什么？\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "b9326e04",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/02-intermediate/recurrent_neural_network/main.py \n",
    "# http://colah.github.io/posts/2015-08-Understanding-LSTMs/\n",
    "# 参考：https://zhuanlan.zhihu.com/p/79064602\n",
    "import torch \n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from tqdm import tqdm\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0f861219",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 使用LSTM的思路\n",
    "# 把图片看成一个个句子，图片的尺寸为28*28，图片的长28相当于28个单词，图片的宽28相当于每个单词的向量长度\n",
    "sequence_length = 28\n",
    "input_size = 28\n",
    "hidden_size = 128\n",
    "num_layers = 2\n",
    "num_classes = 10\n",
    "batch_size = 100\n",
    "num_epochs = 30\n",
    "learning_rate = 0.01\n",
    "gamma = 0.9\n",
    "step_size=5\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "0ece489f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.36it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 77.71it/s]\n",
      "  1%|          | 6/938 [00:00<00:15, 58.28it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 0, train loss:0.6564, train acc:0.7587\n",
      "test acc:0.8299\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.25it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 77.84it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 53.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 1, train loss:0.4203, train acc:0.8458\n",
      "test acc:0.8494\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 55.71it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.37it/s]\n",
      "  1%|          | 6/938 [00:00<00:15, 58.96it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 2, train loss:0.3625, train acc:0.8664\n",
      "test acc:0.8565\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 55.92it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.03it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 54.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 3, train loss:0.3302, train acc:0.8774\n",
      "test acc:0.8729\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.09it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.66it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 52.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 4, train loss:0.3012, train acc:0.8884\n",
      "test acc:0.8775\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.05it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.32it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 54.81it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 5, train loss:0.2803, train acc:0.8960\n",
      "test acc:0.8844\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.34it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.08it/s]\n",
      "  1%|          | 6/938 [00:00<00:18, 51.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 6, train loss:0.2610, train acc:0.9027\n",
      "test acc:0.8853\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.19it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 76.39it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 55.76it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 7, train loss:0.2401, train acc:0.9115\n",
      "test acc:0.8874\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.27it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.19it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 53.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 8, train loss:0.2197, train acc:0.9192\n",
      "test acc:0.8925\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.08it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.17it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 56.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 9, train loss:0.2030, train acc:0.9250\n",
      "test acc:0.8933\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.08it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 77.04it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 53.96it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:10, train loss:0.1806, train acc:0.9321\n",
      "test acc:0.8999\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.26it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.60it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 52.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:11, train loss:0.1614, train acc:0.9399\n",
      "test acc:0.8984\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.15it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.61it/s]\n",
      "  1%|          | 6/938 [00:00<00:18, 50.86it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:12, train loss:0.1447, train acc:0.9461\n",
      "test acc:0.9006\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.02it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.10it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 53.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:13, train loss:0.1255, train acc:0.9540\n",
      "test acc:0.9020\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.08it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.86it/s]\n",
      "  1%|          | 5/938 [00:00<00:19, 46.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:14, train loss:0.1085, train acc:0.9601\n",
      "test acc:0.9030\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 55.88it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.07it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 54.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:15, train loss:0.0916, train acc:0.9674\n",
      "test acc:0.9032\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.32it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.68it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 56.06it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:16, train loss:0.0751, train acc:0.9737\n",
      "test acc:0.8995\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.24it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.40it/s]\n",
      "  1%|          | 6/938 [00:00<00:15, 59.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:17, train loss:0.0628, train acc:0.9788\n",
      "test acc:0.9005\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.28it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.63it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 56.57it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:18, train loss:0.0507, train acc:0.9841\n",
      "test acc:0.9031\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.32it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.80it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 55.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:19, train loss:0.0409, train acc:0.9878\n",
      "test acc:0.9034\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.45it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.00it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 58.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:20, train loss:0.0315, train acc:0.9913\n",
      "test acc:0.9025\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.47it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 79.39it/s]\n",
      "  1%|          | 7/938 [00:00<00:15, 60.37it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:21, train loss:0.0251, train acc:0.9933\n",
      "test acc:0.9037\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.64it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.96it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 56.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:22, train loss:0.0186, train acc:0.9956\n",
      "test acc:0.9028\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.58it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 80.02it/s]\n",
      "  1%|          | 6/938 [00:00<00:18, 51.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:23, train loss:0.0138, train acc:0.9971\n",
      "test acc:0.9024\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.59it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 80.29it/s]\n",
      "  1%|          | 5/938 [00:00<00:18, 49.56it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:24, train loss:0.0101, train acc:0.9983\n",
      "test acc:0.9003\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.10it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.54it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 56.30it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:25, train loss:0.0090, train acc:0.9986\n",
      "test acc:0.8999\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.19it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 77.57it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 53.78it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:26, train loss:0.0059, train acc:0.9992\n",
      "test acc:0.9019\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 55.90it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.36it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 55.02it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:27, train loss:0.0057, train acc:0.9991\n",
      "test acc:0.9014\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 55.94it/s]\n",
      "100%|██████████| 157/157 [00:02<00:00, 78.20it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 52.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:28, train loss:0.0040, train acc:0.9995\n",
      "test acc:0.9011\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 56.03it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 78.81it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:29, train loss:0.0033, train acc:0.9997\n",
      "test acc:0.9000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "class RNN(nn.Module):\n",
    "    def __init__(self, input_size, hidden_size, num_layers, num_classes):\n",
    "        super(RNN, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_layers = num_layers\n",
    "        # 官网默认的输入是seq_len, batch, input_size,这里将batch first=True放在第一个\n",
    "        # seq len 相当于一句话的长度， input size相当于每个单词的向量长度\n",
    "        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True)\n",
    "        self.fc = nn.Linear(hidden_size, num_classes)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # 初始化两个向量的值, 因为这里是两层的lstm\n",
    "        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device) \n",
    "        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(device)\n",
    "        \n",
    "        # 前向传播，这里只需要out即可\n",
    "        out, _ = self.lstm(x, (h0, c0))  # out: tensor of shape (batch_size, seq_length, hidden_size)\n",
    "        \n",
    "        # 使用全连层进行分类\n",
    "        out = self.fc(out[:, -1, :])\n",
    "        return out\n",
    "    \n",
    "model = RNN(input_size, hidden_size, num_layers, num_classes).to(device)\n",
    "\n",
    "\n",
    "# Loss and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma) # 学习方式\n",
    "# Train the model\n",
    "total_step = len(train_loader)\n",
    "train_acc, test_acc = [], []\n",
    "for epoch in range(num_epochs):\n",
    "    epoch_accuracy = 0\n",
    "    epoch_loss = 0\n",
    "    model.train()\n",
    "    for i, (images, labels) in enumerate(tqdm(train_loader)):\n",
    "        # reshape的维度变成了(batch size, 28, 28)\n",
    "        images = images.reshape(-1, sequence_length, input_size).to(device)\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        # Backward and optimize\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        acc = (outputs.argmax(dim=1) == labels).float().mean()\n",
    "        epoch_accuracy += acc / len(train_loader)\n",
    "        epoch_loss += loss / len(train_loader)\n",
    "        \n",
    "    # Test the model\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        epoch_test_accuracy = 0\n",
    "        epoch_test_loss = 0\n",
    "        for images, labels in tqdm(test_loader):\n",
    "            images = images.reshape(-1, sequence_length, input_size).to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            acc = (predicted == labels).float().mean()\n",
    "            epoch_test_accuracy += acc / len(test_loader)\n",
    "\n",
    "    scheduler.step()\n",
    "    print(f'EPOCH:{epoch:2}, train loss:{epoch_loss:.4f}, train acc:{epoch_accuracy:.4f}')\n",
    "    print(f'test acc:{epoch_test_accuracy:.4f}')\n",
    "\n",
    "    train_acc.append(epoch_accuracy)\n",
    "    test_acc.append(epoch_test_accuracy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f0e9611",
   "metadata": {},
   "source": [
    "# Transformer网络"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "8b22031d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://bitjoy.net/2019/08/02/cs224n%ef%bc%881-31%ef%bc%89translation-seq2seq-attention/\n",
    "# https://bitjoy.net/2020/03/04/cs224n%ef%bc%882-21%ef%bc%89transformers-and-self-attention-for-generative-models/\n",
    "# https://zhuanlan.zhihu.com/p/48731949\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import TensorDataset, DataLoader, Dataset\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "10e5ee2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model=28, max_len=28):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        pe = torch.zeros(max_len, d_model)\n",
    "        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))\n",
    "        pe[:, 0::2] = torch.sin(position * div_term)\n",
    "        pe[:, 1::2] = torch.cos(position * div_term)\n",
    "        pe = pe.unsqueeze(0).transpose(0, 1)\n",
    "        self.register_buffer('pe', pe)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.pe[:x.size(0), :]        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "id": "69ade7ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    # 这里feature size必须可以整除以nhead\n",
    "    def __init__(self, feature_size=28, num_layers=3, nhead=7):\n",
    "        super().__init__()\n",
    "        self.src_mask = None\n",
    "        self.pos_encoder = PositionalEncoding(feature_size).to(device)\n",
    "        self.nhead = nhead\n",
    "        self.feature_size = feature_size\n",
    "        self.encoder_layer = nn.TransformerEncoderLayer(d_model=feature_size, nhead=self.nhead)\n",
    "        self.transformer_encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=num_layers)\n",
    "        self.decoder = nn.Linear(feature_size**2, 10)\n",
    "        self.init_weight()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        if self.src_mask is None:\n",
    "            mask = self._generate_square_subsequent_mask(len(x)).to(device)\n",
    "            self.src_mask = mask\n",
    "        x = self.pos_encoder(x)    \n",
    "        x = self.transformer_encoder(x, self.src_mask)\n",
    "        x = self.decoder(x.permute(1, 0, 2).reshape(-1, self.feature_size**2))\n",
    "        return x\n",
    "    \n",
    "   # 将decoder 的偏置转换成0， 将权重进行标准化 \n",
    "    def init_weight(self):\n",
    "        self.decoder.bias.data.zero_()\n",
    "        \n",
    "   # 将输入的序列进行mask操作 \n",
    "    def _generate_square_subsequent_mask(self, sz):\n",
    "        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)\n",
    "        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
    "        return mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "id": "66470857",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = Transformer().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "id": "da202308",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 100\n",
    "num_epochs = 30\n",
    "learning_rate = 0.001\n",
    "gamma = 0.9\n",
    "step_size=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "id": "0eea378a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.22it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.43it/s]\n",
      "  1%|          | 5/938 [00:00<00:20, 45.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 0, train loss:0.4778, train acc:0.8260\n",
      "test acc:0.8542\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.33it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.72it/s]\n",
      "  1%|          | 5/938 [00:00<00:21, 42.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 1, train loss:0.3591, train acc:0.8700\n",
      "test acc:0.8720\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.03it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 100.89it/s]\n",
      "  1%|          | 5/938 [00:00<00:21, 43.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 2, train loss:0.3222, train acc:0.8818\n",
      "test acc:0.8797\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.82it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 101.87it/s]\n",
      "  1%|          | 5/938 [00:00<00:20, 45.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 3, train loss:0.2953, train acc:0.8914\n",
      "test acc:0.8799\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 46.95it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 101.78it/s]\n",
      "  1%|          | 5/938 [00:00<00:22, 41.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 4, train loss:0.2744, train acc:0.8985\n",
      "test acc:0.8927\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.16it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.75it/s]\n",
      "  1%|          | 5/938 [00:00<00:21, 43.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 5, train loss:0.2569, train acc:0.9055\n",
      "test acc:0.8931\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.30it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.52it/s]\n",
      "  1%|          | 5/938 [00:00<00:20, 44.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 6, train loss:0.2430, train acc:0.9089\n",
      "test acc:0.8947\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.34it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 104.83it/s]\n",
      "  1%|          | 5/938 [00:00<00:20, 44.83it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 7, train loss:0.2258, train acc:0.9168\n",
      "test acc:0.8953\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:18<00:00, 51.43it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.93it/s]\n",
      "  1%|          | 6/938 [00:00<00:18, 51.73it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 8, train loss:0.2151, train acc:0.9193\n",
      "test acc:0.8997\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 46.99it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.75it/s]\n",
      "  0%|          | 0/938 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH: 9, train loss:0.2031, train acc:0.9249\n",
      "test acc:0.8997\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.78it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.46it/s]\n",
      "  1%|          | 6/938 [00:00<00:16, 55.68it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:10, train loss:0.1911, train acc:0.9294\n",
      "test acc:0.8982\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 48.48it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 101.37it/s]\n",
      "  1%|          | 5/938 [00:00<00:22, 40.94it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:11, train loss:0.1803, train acc:0.9329\n",
      "test acc:0.9027\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.17it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 100.43it/s]\n",
      "  0%|          | 0/938 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:12, train loss:0.1703, train acc:0.9367\n",
      "test acc:0.9064\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.89it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 101.35it/s]\n",
      "  0%|          | 0/938 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:13, train loss:0.1604, train acc:0.9394\n",
      "test acc:0.9078\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 46.93it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.09it/s]\n",
      "  1%|          | 5/938 [00:00<00:21, 43.51it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:14, train loss:0.1534, train acc:0.9438\n",
      "test acc:0.9062\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.06it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.23it/s]\n",
      "  0%|          | 0/938 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:15, train loss:0.1452, train acc:0.9461\n",
      "test acc:0.9051\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.31it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 101.36it/s]\n",
      "  1%|          | 5/938 [00:00<00:22, 40.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:16, train loss:0.1387, train acc:0.9486\n",
      "test acc:0.9045\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:18<00:00, 49.73it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 101.86it/s]\n",
      "  0%|          | 0/938 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:17, train loss:0.1332, train acc:0.9512\n",
      "test acc:0.9042\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 49.32it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.58it/s]\n",
      "  1%|          | 5/938 [00:00<00:21, 43.66it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:18, train loss:0.1264, train acc:0.9540\n",
      "test acc:0.9071\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.13it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.38it/s]\n",
      "  1%|          | 5/938 [00:00<00:20, 46.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:19, train loss:0.1190, train acc:0.9560\n",
      "test acc:0.9078\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.04it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 98.46it/s]\n",
      "  1%|          | 5/938 [00:00<00:22, 41.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:20, train loss:0.1160, train acc:0.9573\n",
      "test acc:0.9064\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 45.90it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 98.39it/s] \n",
      "  1%|          | 5/938 [00:00<00:20, 45.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:21, train loss:0.1106, train acc:0.9589\n",
      "test acc:0.9064\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.21it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.20it/s]\n",
      "  1%|          | 5/938 [00:00<00:20, 44.96it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:22, train loss:0.1078, train acc:0.9604\n",
      "test acc:0.9084\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.88it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.83it/s]\n",
      "  1%|          | 5/938 [00:00<00:21, 42.61it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:23, train loss:0.1040, train acc:0.9611\n",
      "test acc:0.9061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.27it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.54it/s]\n",
      "  1%|          | 5/938 [00:00<00:22, 41.18it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:24, train loss:0.1011, train acc:0.9629\n",
      "test acc:0.9070\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.82it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 100.92it/s]\n",
      "  1%|          | 5/938 [00:00<00:18, 49.85it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:25, train loss:0.0975, train acc:0.9633\n",
      "test acc:0.9087\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:17<00:00, 54.72it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 103.59it/s]\n",
      "  1%|          | 6/938 [00:00<00:17, 54.79it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:26, train loss:0.0962, train acc:0.9649\n",
      "test acc:0.9092\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:16<00:00, 55.58it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 100.72it/s]\n",
      "  1%|          | 7/938 [00:00<00:15, 61.36it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:27, train loss:0.0925, train acc:0.9668\n",
      "test acc:0.9083\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:19<00:00, 47.29it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 102.32it/s]\n",
      "  1%|          | 5/938 [00:00<00:21, 43.95it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:28, train loss:0.0893, train acc:0.9675\n",
      "test acc:0.9085\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 938/938 [00:20<00:00, 46.45it/s]\n",
      "100%|██████████| 157/157 [00:01<00:00, 101.87it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "EPOCH:29, train loss:0.0882, train acc:0.9676\n",
      "test acc:0.9081\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Loss and optimizer\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma) # 学习方式\n",
    "# Train the model\n",
    "total_step = len(train_loader)\n",
    "train_acc, test_acc = [], []\n",
    "# seq_len 相当于 max len， 即每句话切分的最大长度\n",
    "sequence_length = 28\n",
    "# input_size 相当于embedding dim\n",
    "input_size = 28\n",
    "for epoch in range(num_epochs):\n",
    "    epoch_accuracy = 0\n",
    "    epoch_loss = 0\n",
    "    model.train()\n",
    "    for i, (images, labels) in enumerate(tqdm(train_loader)):\n",
    "        # reshape的维度变成了(batch size, 28, 28)\n",
    "        images = images.reshape(-1, sequence_length, input_size)\n",
    "        images = images.permute(1, 0, 2).to(device)\n",
    "        labels = labels.to(device)\n",
    "        \n",
    "        # Forward pass\n",
    "        outputs = model(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        \n",
    "        # Backward and optimize\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        acc = (outputs.argmax(dim=1) == labels).float().mean()\n",
    "        epoch_accuracy += acc / len(train_loader)\n",
    "        epoch_loss += loss / len(train_loader)\n",
    "        \n",
    "    # Test the model\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        epoch_test_accuracy = 0\n",
    "        epoch_test_loss = 0\n",
    "        for images, labels in tqdm(test_loader):\n",
    "            images = images.reshape(-1, sequence_length, input_size)\n",
    "            images = images.permute(1, 0, 2).to(device)\n",
    "            labels = labels.to(device)\n",
    "            outputs = model(images)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            acc = (predicted == labels).float().mean()\n",
    "            epoch_test_accuracy += acc / len(test_loader)\n",
    "\n",
    "    scheduler.step()\n",
    "    print(f'EPOCH:{epoch:2}, train loss:{epoch_loss:.4f}, train acc:{epoch_accuracy:.4f}')\n",
    "    print(f'test acc:{epoch_test_accuracy:.4f}')\n",
    "\n",
    "    train_acc.append(epoch_accuracy)\n",
    "    test_acc.append(epoch_test_accuracy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "3d0f23b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "9a59465a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Transformer().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "ec9a2ddd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([28, 64, 10])"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(d).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "ad66feca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3, 2, 2])"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.randn(3, 2).unsqueeze(1).expand(-1, 2, -1).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca0e0e09",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "173px"
   },
   "toc_section_display": true,
   "toc_window_display": true
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
